package proxy

import (
	"fmt"
	"net/http"
	"sync"
)

var _ ModelServiceConnPool = (*RoundRobinConnPool)(nil)

type Client struct {
	HttpClient *http.Client
	Addr       string
	Available  bool
}

type RoundRobinConnPool struct {
	Clients []*Client
	current int
	sync.Mutex
}

func NewRoundRobinConnPool(baseURLs []string) *RoundRobinConnPool {
	clients := make([]*Client, len(baseURLs))
	for i, addr := range baseURLs {
		clients[i] = &Client{
			HttpClient: &http.Client{},
			Addr:       addr,
			Available:  true,
		}
	}
	return &RoundRobinConnPool{
		Clients: clients,
		current: 0,
	}
}

func (p *RoundRobinConnPool) GetOneAvailableConn() (*http.Client, error) {
	p.Lock()
	defer p.Unlock()

	for i := 0; i < len(p.Clients); i++ {
		client := p.Clients[p.current]
		p.current = (p.current + 1) % len(p.Clients)
		if client.Available {
			return client.HttpClient, nil
		}
	}
	return nil, fmt.Errorf("no available connection")
}

func (p *RoundRobinConnPool) MarkConnUnavailable(addr string) {
	p.Lock()
	defer p.Unlock()

	for i := range p.Clients {
		if p.Clients[i].Addr == addr {
			p.Clients[i].Available = false
			return
		}
	}
}

func (p *RoundRobinConnPool) MarkConnAvailable(addr string) {
	p.Lock()
	defer p.Unlock()

	for i := range p.Clients {
		if p.Clients[i].Addr == addr {
			p.Clients[i].Available = true
			return
		}
	}
}

func (p *RoundRobinConnPool) HealthCheck() {
	for _, client := range p.Clients {
		go func(c *Client) {
			resp, err := c.HttpClient.Get(c.Addr + "/healthz")
			if err != nil || resp.StatusCode != http.StatusOK {
				p.MarkConnUnavailable(c.Addr)
				return
			}
			defer resp.Body.Close()
			p.MarkConnAvailable(c.Addr)
		}(client)
	}
}
